#include <ntddk.h>

#include "..\ZeroCommon\ZeroCommon.h"

#define DRIVER_PREFIX "Zero: "

long long g_TotalRead, g_TotalWrite;

VOID
ZeroUnload(_In_ PDRIVER_OBJECT DriverObject);

NTSTATUS
ZeroCreateClose(_In_ PDEVICE_OBJECT DeivceObject, _In_ PIRP Irp);

NTSTATUS
ZeroRead(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp);

NTSTATUS
ZeroWrite(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp);

NTSTATUS
ZeroDeviceControl(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp);

NTSTATUS
CompleteIrp(PIRP Irp, NTSTATUS Status = STATUS_SUCCESS, ULONG_PTR Info = 0);

EXTERN_C
NTSTATUS
DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegisterPath)
{
	UNREFERENCED_PARAMETER(RegisterPath);

	DriverObject->DriverUnload                         = ZeroUnload;

	DriverObject->MajorFunction[IRP_MJ_CREATE]         = ZeroCreateClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE]          = ZeroCreateClose;
	DriverObject->MajorFunction[IRP_MJ_READ]           = ZeroRead;
	DriverObject->MajorFunction[IRP_MJ_WRITE]          = ZeroWrite;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = ZeroDeviceControl;

	UNICODE_STRING DevName          = RTL_CONSTANT_STRING(L"\\Device\\Zero");
	UNICODE_STRING SymLink          = RTL_CONSTANT_STRING(L"\\??\\Zero");
	PDEVICE_OBJECT DeviveObject     = nullptr;
	NTSTATUS Status                 = STATUS_SUCCESS;

	do {
		Status = IoCreateDevice(DriverObject, 0, &DevName, FILE_DEVICE_UNKNOWN, 0, FALSE, &DeviveObject);
		if (!NT_SUCCESS(Status)) {
			KdPrint((DRIVER_PREFIX "falied to create device (0x%08X)\n", Status));
			break;
		}

		DeviveObject->Flags |= DO_DIRECT_IO;

		Status = IoCreateSymbolicLink(&SymLink, &DevName);
		if (!NT_SUCCESS(Status)) {
			KdPrint((DRIVER_PREFIX "failed to create symbolic link (0x%08X)\n", Status));
			break;
		}
		KdPrint((DRIVER_PREFIX "Driver loaded successfully\n"));
	} while (FALSE);

	if (!NT_SUCCESS(Status)) {
		if (DeviveObject != nullptr) {
			IoDeleteDevice(DeviveObject);
		}
		KdPrint((DRIVER_PREFIX "Driver loaded error\n"));
	}

	return Status;
}

VOID
ZeroUnload(_In_ PDRIVER_OBJECT DriverObject)
{
	UNICODE_STRING SymLink = RTL_CONSTANT_STRING(L"\\??\\Zero");
	IoDeleteSymbolicLink(&SymLink);

	IoDeleteDevice(DriverObject->DeviceObject);

	KdPrint((DRIVER_PREFIX "Driver uninstalled successfully\n"));
}

NTSTATUS
ZeroCreateClose(_In_ PDEVICE_OBJECT DeivceObject, _In_ PIRP Irp)
{
	return CompleteIrp(Irp);
}

NTSTATUS
ZeroRead(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp)
{
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	ULONG Len                = Stack->Parameters.Read.Length;
	if (Len == 0) {
		return CompleteIrp(Irp, STATUS_INVALID_BUFFER_SIZE);
	}

	PVOID Buffer = MmGetSystemAddressForMdlSafe(Irp->MdlAddress, NormalPagePriority);
	if (Buffer == nullptr) {
		return CompleteIrp(Irp, STATUS_INSUFFICIENT_RESOURCES);
	}

	memset(Buffer, 0, Len);

	InterlockedAdd64(&g_TotalRead, Len);

	return CompleteIrp(Irp, STATUS_SUCCESS, Len);
}

NTSTATUS
ZeroWrite(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp)
{
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	ULONG Len                = Stack->Parameters.Write.Length;

	InterlockedAdd64(&g_TotalWrite, Len);

	return CompleteIrp(Irp, STATUS_SUCCESS, Len);
}

NTSTATUS
ZeroDeviceControl(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp)
{
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	
	if (Stack->Parameters.DeviceIoControl.IoControlCode != IOCTL_ZERO_GET_STATS) {
		return CompleteIrp(Irp, STATUS_INVALID_DEVICE_REQUEST);
	}

	if (Stack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(ZeroStatus)) {
		return CompleteIrp(Irp, STATUS_BUFFER_TOO_SMALL);
	}

	ZeroStatus* Status = (ZeroStatus*)(Irp->AssociatedIrp.SystemBuffer);
	Status->TotalRead  = g_TotalRead;
	Status->TotalWrite = g_TotalWrite;

	return CompleteIrp(Irp, STATUS_SUCCESS, sizeof(Status));
}

NTSTATUS
CompleteIrp(PIRP Irp, NTSTATUS Status, ULONG_PTR Info)
{
	Irp->IoStatus.Status      = Status;
	Irp->IoStatus.Information = Info;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Status;
}